package com.han.agent.interview;

import io.micrometer.observation.ObservationRegistry;
import org.springframework.ai.embedding.observation.EmbeddingModelObservationConvention;
import org.springframework.ai.model.SimpleApiKey;
import org.springframework.ai.model.openai.autoconfigure.OpenAIAutoConfigurationUtil;
import org.springframework.ai.model.openai.autoconfigure.OpenAiChatProperties;
import org.springframework.ai.model.openai.autoconfigure.OpenAiConnectionProperties;
import org.springframework.ai.model.openai.autoconfigure.OpenAiEmbeddingAutoConfiguration;
import org.springframework.ai.model.openai.autoconfigure.OpenAiEmbeddingProperties;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;

import static org.springframework.ai.model.openai.autoconfigure.OpenAIAutoConfigurationUtil.resolveConnectionProperties;

/**
 * {@link OpenAiEmbeddingAutoConfiguration}
 *
 * @author han
 */
@Configuration
public class Config {

    @Bean
    public OpenAiEmbeddingModel openAiEmbeddingModel(OpenAiConnectionProperties commonProperties,
                                                     OpenAiEmbeddingProperties embeddingProperties, ObjectProvider<RestClient.Builder> restClientBuilderProvider,
                                                     ObjectProvider<WebClient.Builder> webClientBuilderProvider, RetryTemplate retryTemplate,
                                                     ResponseErrorHandler responseErrorHandler, ObjectProvider<ObservationRegistry> observationRegistry,
                                                     ObjectProvider<EmbeddingModelObservationConvention> observationConvention) {

        var openAiApi = openAiApi(embeddingProperties, commonProperties,
                restClientBuilderProvider.getIfAvailable(RestClient::builder),
                webClientBuilderProvider.getIfAvailable(WebClient::builder), responseErrorHandler, "embedding");

        var embeddingModel = new CustomOpenAiEmbeddingModel(openAiApi, embeddingProperties.getMetadataMode(),
                embeddingProperties.getOptions(), retryTemplate,
                observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));

        observationConvention.ifAvailable(embeddingModel::setObservationConvention);

        return embeddingModel;
    }

    private OpenAiApi openAiApi(OpenAiEmbeddingProperties embeddingProperties,
                                OpenAiConnectionProperties commonProperties, RestClient.Builder restClientBuilder,
                                WebClient.Builder webClientBuilder, ResponseErrorHandler responseErrorHandler, String modelType) {

        OpenAIAutoConfigurationUtil.ResolvedConnectionProperties resolved = resolveConnectionProperties(
                commonProperties, embeddingProperties, modelType);

        return OpenAiApi.builder()
                .baseUrl(resolved.baseUrl())
                .apiKey(new SimpleApiKey(resolved.apiKey()))
                .headers(resolved.headers())
                .completionsPath(OpenAiChatProperties.DEFAULT_COMPLETIONS_PATH)
                .embeddingsPath(embeddingProperties.getEmbeddingsPath())
                .restClientBuilder(restClientBuilder)
                .webClientBuilder(webClientBuilder)
                .responseErrorHandler(responseErrorHandler)
                .build();
    }

}
